import os
import tkinter as tk
from tkinter.filedialog import askopenfilename
from PIL import Image, ImageTk
import torch
from torchvision import transforms
# from model import AlexNet
from networks.ClassicNetwork.ResNet import ResNet50
from networks.ClassicNetwork.VGGNet import VGG16
from networks.ClassicNetwork.AlexNet import AlexNet
class MainWindow:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title('识别')

        self.frame_left = tk.Frame(self.root, padx=1, pady=5, bg="#aaaaaa")
        self.frame_left.pack(padx=5, pady=10, fill="y", side=tk.LEFT)

        self.frame_right = tk.Frame(self.root, padx=5, pady=5)
        self.frame_right.pack(padx=5, pady=10, fill="y", side=tk.LEFT)


        self.button_loadMode_l = tk.Button(self.frame_right, text='加载VggNet模型', command=self.loadresult_vgg, width=15, height=2)
        self.button_loadMode_l.pack(fill="x")

        self.button_loadMode_l = tk.Button(self.frame_right, text='加载AlexNet模型', command=self.loadModel_ale, width=15, height=2)
        self.button_loadMode_l.pack(fill="x")

        self.button_loadModel = tk.Button(self.frame_right, text='加载Resnet模型', command=self.loadModel, width=15, height=2)
        self.button_loadModel.pack(fill="x")

        self.button_loadImage = tk.Button(self.frame_right, text='加载图像', command=self.loadimg, width=15, height=2)
        self.button_loadImage.pack(fill="x")
        self.button_loadImage.config(state=tk.DISABLED)

        self.button_predict = tk.Button(self.frame_right, text='识别一下', command=self.predict, width=15, height=2)
        self.button_predict.pack(fill="x")
        self.button_predict.config(state=tk.DISABLED)

        self.label_info = tk.Label(self.frame_right, font=('宋体', 12), justify=tk.LEFT, padx=2, pady=30)
        self.label_info.pack(fill="x")
        self.label_info.config(text="请载入一个模型文件")

        self.canvas = tk.Canvas(self.frame_left, bg='gray', height=450, width=450)
        self.canvas.pack(fill='x', expand='yes')

        self.frame_left_1 = tk.Frame(self.root, padx=1, pady=5, bg="#aaaaaa")
        self.frame_left_1.pack(padx=5, pady=10, fill="y", side=tk.LEFT)
        self.canvas1 = tk.Canvas(self.frame_left_1, bg='gray', height=250, width=250)
        self.canvas1.pack(fill='x', expand='yes')
        self.canvas4 = tk.Canvas(self.frame_left_1, bg='gray', height=250, width=250)
        self.canvas4.pack(fill='x', expand='yes')

        self.frame_left_2 = tk.Frame(self.root, padx=1, pady=5, bg="#aaaaaa")
        self.frame_left_2.pack(padx=5, pady=10, fill="y", side=tk.BOTTOM)

        self.canvas2 = tk.Canvas(self.frame_left_2, bg='gray', height=250, width=250)
        self.canvas2.pack(fill='x', expand='yes')

        self.canvas3 = tk.Canvas(self.frame_left_2, bg='gray', height=250, width=250)
        self.canvas3.pack(fill='x', expand='yes')

        self.root.mainloop()



    def load_image_to_canvas(self, file_path):
        """把给定路径的图像加载入self.img 并绘制到canvas"""

        def resize(w_box, h_box, pil_image):  # 参数是：要适应的窗口宽、高、Image.open后的图片
            w, h = pil_image.size  # 获取图像的原始大小
            f1 = 1.0 * w_box / w
            f2 = 1.0 * h_box / h
            factor = min([f1, f2])
            width = int(w * factor)
            height = int(h * factor)
            return pil_image.resize((width, height), Image.ANTIALIAS)

        try:
            img = Image.open(file_path)
            self.img = img
            img_w, img_h = 500,500
            img = resize(img_w, img_h, img)
            self.pil_img = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
            self.canvas.update()  # 获取宽高之前要先对于这个组件update()
            x, y = 0, (self.canvas.winfo_height() - img.size[1]) / 2
            self.canvas.create_image(x, y, anchor='nw', image=self.pil_img)
        except Exception as e:
            self.label_info.config(text="图片载入出错")
        finally:
            self.button_predict.config(state=tk.NORMAL)
            self.label_info.config(text="图片已载入\n点击预测按钮")

    def load_image_to_canvas1(self, file_path):
        """把给定路径的图像加载入self.img 并绘制到canvas"""

        def resize(w_box, h_box, pil_image):  # 参数是：要适应的窗口宽、高、Image.open后的图片
            w, h = pil_image.size  # 获取图像的原始大小
            f1 = 1.0 * w_box / w
            f2 = 1.0 * h_box / h
            factor = min([f1, f2])
            width = int(w * factor)
            height = int(h * factor)
            return pil_image.resize((width, height), Image.ANTIALIAS)

        try:
            img = Image.open(file_path)
            self.img = img
            img_w, img_h = img.size
            if img_w > 400:
                img_w = 400
                img_h = img_h * (400 / img_w)
                img = resize(img_w, img_h, img)
            self.pil_img = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
            self.canvas1.update()  # 获取宽高之前要先对于这个组件update()
            x, y = 0, (self.canvas1.winfo_height() - img.size[1]) / 2
            self.canvas1.create_image(x, y, anchor='nw', image=self.pil_img)
        except Exception as e:
            self.label_info.config(text="图片载入出错")
        finally:
            self.button_predict.config(state=tk.NORMAL)
            self.label_info.config(text="图片已载入\n点击预测按钮")

    def load_image_to_canvas2(self, file_path):
        """把给定路径的图像加载入self.img 并绘制到canvas"""

        def resize(w_box, h_box, pil_image):  # 参数是：要适应的窗口宽、高、Image.open后的图片
            w, h = pil_image.size  # 获取图像的原始大小
            f1 = 1.0 * w_box / w
            f2 = 1.0 * h_box / h
            factor = min([f1, f2])
            width = int(w * factor)
            height = int(h * factor)
            return pil_image.resize((width, height), Image.ANTIALIAS)

        try:
            img = Image.open(file_path)
            self.img = img
            img_w, img_h = img.size
            if img_w > 400:
                img_w = 400
                img_h = img_h * (400 / img_w)
                img = resize(img_w, img_h, img)
            self.pil_img = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
            # self.canvas2.update()  # 获取宽高之前要先对于这个组件update()
            x, y = 0, (self.canvas1.winfo_height() - img.size[1]) / 2
            self.canvas2.create_image(x, y, anchor='nw', image=self.pil_img)
        except Exception as e:
            self.label_info.config(text="图片载入出错")
        finally:
            self.button_predict.config(state=tk.NORMAL)
            self.label_info.config(text="图片已载入\n点击预测按钮")


    def predict(self):
        """根据已载入的模型进行识别"""
        class_dict = {
            "0": "0",
            "1": "1",
        }
        preprocess_transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        img = preprocess_transform(self.img)
        img.unsqueeze_(0)
        print(img.shape)
        with torch.no_grad():
            outputs, features = self.model(img)

            output = torch.squeeze(outputs)
            predict = torch.softmax(output, dim=0)
            predict_cla = torch.argmax(predict).numpy()  # 最大值位置索引
            print(predict_cla)
        class_str = class_dict[str(predict_cla)]
        prob_str = "%.1f" % (predict[predict_cla].item() * 100)
        self.label_info.config(text=f"类别：{class_str}\n可能性：{prob_str}%")




    def loadModel(self):

        """载入指定的模型"""
        def resize(w_box, h_box, pil_image):  # 参数是：要适应的窗口宽、高、Image.open后的图片
            w, h = pil_image.size  # 获取图像的原始大小
            f1 = 1.0 * w_box / w
            f2 = 1.0 * h_box / h
            factor = min([f1, f2])
            width = int(w * factor)
            height = int(h * factor)
            return pil_image.resize((width, height), Image.ANTIALIAS)
        print("load    resnet")
        print("-------end")
        file_path = r"model/AlexNet_leaf/results/Accuracy.jpg"
        img = Image.open(file_path)
        self.img = img
        img_w, img_h = 250, 250
        img = resize(img_w, img_h, img)
        self.pil_img1 = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
        # self.canvas1.update()  # 获取宽高之前要先对于这个组件update()
        x, y = 0, (self.canvas1.winfo_height() - img.size[1]) / 2
        self.canvas1.create_image(x, y, anchor='nw', image=self.pil_img1)

        file_path = r"model/AlexNet_leaf/results/Loss.jpg"
        img = Image.open(file_path)
        self.img = img
        img_w, img_h = 250, 250
        img = resize(img_w, img_h, img)
        self.pil_img2 = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
        # self.canvas2.update()  # 获取宽高之前要先对于这个组件update()
        x, y = 0, (self.canvas2.winfo_height() - img.size[1]) / 2
        self.canvas2.create_image(x, y, anchor='nw', image=self.pil_img2)

        file_path = r"model/AlexNet_leaf/results/cm.jpg"
        img = Image.open(file_path)
        self.img = img
        img_w, img_h = 250, 250
        img = resize(img_w, img_h, img)
        self.pil_img3 = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
        # self.canvas3.update()  # 获取宽高之前要先对于这个组件update()
        x, y = 0, (self.canvas3.winfo_height() - img.size[1]) / 2
        self.canvas3.create_image(x, y, anchor='nw', image=self.pil_img3)

        file_path = r"model/AlexNet_leaf/results/ROC_plot.jpg"
        img = Image.open(file_path)
        self.img = img
        img_w, img_h = 250, 250
        img = resize(img_w, img_h, img)
        self.pil_img4 = ImageTk.PhotoImage(img)  # PhotoImage返回的对象必须一直被引用着，一旦失去引用，canvas上的图像立即消失
        # self.canvas4.update()  # 获取宽高之前要先对于这个组件update()
        x, y = 0, (self.canvas4.winfo_height() - img.size[1]) / 2
        self.canvas4.create_image(x, y, anchor='nw', image=self.pil_img4)

        try:
            default_dir = os.getcwd()
            modelPath = askopenfilename(title='选择一个模型文件',
                                        initialdir=(os.path.expanduser(default_dir)),
                                        filetypes=[('pkl文件', '*.pkl'), ('All Files', '*')])
            if modelPath == "":
                return
            self.label_info.config(text="载入模型中……")
            model = ResNet50(num_classes=2)
            model.load_state_dict(torch.load(modelPath))
            model.eval()
            self.model = model
            """根据已载入的模型进行识别"""
            class_dict = {
                "0": "0",
                "1": "1",
            }
            preprocess_transform = transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
            img = preprocess_transform(self.img)
            img.unsqueeze_(0)
            print(img.shape)
            with torch.no_grad():
                outputs, features = self.model(img)

                output = torch.squeeze(outputs)
                predict = torch.softmax(output, dim=0)
                predict_cla = torch.argmax(predict).numpy()  # 最大值位置索引
                print(predict_cla)
            class_str = class_dict[str(predict_cla)]
            prob_str = "%.1f" % (predict[predict_cla].item() * 100)
            self.label_info.config(text=f"类别：{class_str}\n可能性：{prob_str}%")



        except Exception as e:
            self.label_info.config(text="模型载入出错")
        finally:
            self.button_loadImage.config(state=tk.NORMAL)
            self.label_info.config(text="请打开一张图片")


    def loadresult_vgg(self):
        """载入指定的模型"""
        try:
            default_dir = os.getcwd()
            modelPath = askopenfilename(title='选择一个模型文件',
                                        initialdir=(os.path.expanduser(default_dir)),
                                        filetypes=[('pkl文件', '*.pkl'), ('All Files', '*')])
            if modelPath == "":
                return
            self.label_info.config(text="载入模型中……")
            model = VGG16()
            model.load_state_dict(torch.load(modelPath))
            model.eval()
            self.model = model
        except Exception as e:
            self.label_info.config(text="模型载入出错")
        finally:
            self.button_loadImage.config(state=tk.NORMAL)
            self.label_info.config(text="请打开一张图片")

    def loadModel_ale(self):
        """载入指定的模型"""
        try:
            default_dir = os.getcwd()
            modelPath = askopenfilename(title='选择一个模型文件',
                                        initialdir=(os.path.expanduser(default_dir)),
                                        filetypes=[('pkl文件', '*.pkl'), ('All Files', '*')])
            if modelPath == "":
                return
            self.label_info.config(text="载入模型中……")
            model = AlexNet()
            model.load_state_dict(torch.load(modelPath))
            model.eval()
            self.model = model
        except Exception as e:
            self.label_info.config(text="模型载入出错")
        finally:
            self.button_loadImage.config(state=tk.NORMAL)
            self.label_info.config(text="请打开一张图片")

    def loadimg(self):
        """载入指定的jpg图片"""
        default_dir = os.getcwd()
        photoPath = askopenfilename(title='打开一个照片（png格式）',
                                    initialdir=(os.path.expanduser(default_dir)),
                                    filetypes=[('.png文件', '*.png'), ('All Files', '*')])
        if photoPath == "":
            return
        self.load_image_to_canvas(photoPath)

    def loadresult(self):
        """载入指定的jpg图片"""
        default_dir = os.getcwd()
        photoPath1=  r""
        photoPath2 = r""
        self.load_image_to_canvas1(photoPath1)
        self.load_image_to_canvas2(photoPath2)
if __name__ == '__main__':
    win = MainWindow()